package com.briup.pai.service.impl;

import cn.hutool.core.util.ObjectUtil;
import com.alibaba.fastjson2.JSON;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.briup.pai.common.constant.CommonConstant;
import com.briup.pai.common.constant.ModelConstant;
import com.briup.pai.common.enums.DatasetStatusEnum;
import com.briup.pai.common.enums.ModelOperationTypeEnum;
import com.briup.pai.common.enums.ResultCodeEnum;
import com.briup.pai.common.enums.TrainingStatusEnum;
import com.briup.pai.common.exception.BriupAssert;
import com.briup.pai.common.utils.SecurityUtil;
import com.briup.pai.convert.DatasetConvert;
import com.briup.pai.convert.ModelConfigConvert;
import com.briup.pai.convert.ModelConvert;
import com.briup.pai.entity.dto.ModelOperationDTO;
import com.briup.pai.entity.message.ModelEvaluateMessage;
import com.briup.pai.entity.message.ModelInitTrainMessage;
import com.briup.pai.entity.message.ModelOptimizeTrainMessage;
import com.briup.pai.entity.message.ModelReleaseMessage;
import com.briup.pai.entity.po.*;
import com.briup.pai.entity.vo.ReleaseModelVO;
import com.briup.pai.entity.vo.TrainingDatasetQueryVO;
import com.briup.pai.service.*;
import lombok.SneakyThrows;
import org.springframework.amqp.rabbit.core.RabbitTemplate;
import org.springframework.aop.framework.AopContext;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.http.*;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.multipart.MultipartFile;

import java.util.*;

@Service
public class ModelOperationServiceImpl implements IModelOperationService {

    @Value("${upload.nginx-file-path}")
    private String nginxFilePath;
    @Value("${mq.init.exchange}")
    private String initExchange;
    @Value("${mq.init.send-routing-key}")
    private String initRoutingKey;
    @Value("${mq.optimize.exchange}")
    private String optimizeExchange;
    @Value("${mq.optimize.send-routing-key}")
    private String optimizeRoutingKey;
    @Value("${mq.evaluate.exchange}")
    private String evaluateExchange;
    @Value("${mq.evaluate.send-routing-key}")
    private String evaluateRoutingKey;
    @Value("${mq.release.exchange}")
    private String releaseExchange;
    @Value("${mq.release.send-routing-key}")
    private String releaseRoutingKey;

    @Autowired
    private IModelService modelService;
    @Autowired
    private IDatasetService datasetService;
    @Autowired
    private IModelConfigService modelConfigService;
    @Autowired
    private ITrainingService trainingService;
    @Autowired
    private ITrainingDatasetService trainingDatasetService;
    @Autowired
    private IClassifyService classifyService;
    @Autowired
    private IEntityService entityService;
    @Autowired
    private IEvaluateService evaluateService;
    @Autowired
    private IReleaseService releaseService;

    @Autowired
    private DatasetConvert datasetConvert;
    @Autowired
    private ModelConfigConvert modelConfigConvert;
    @Autowired
    private ModelConvert modelConvert;

    @Autowired
    private RabbitTemplate rabbitTemplate;
    @Autowired
    private RestTemplate restTemplate;

    @Override
    public List<TrainingDatasetQueryVO> getTrainingDatasets(Integer modelId, Integer datasetUsage) {
        // 模型必须存在
        Model model = BriupAssert.requireNotNull(
                modelService,
                Model::getId,
                modelId,
                ResultCodeEnum.DATA_NOT_EXIST
        );
        // 数据集的用途必须和模型一致，数据集得上传完图片了（状态是已完成）
        LambdaQueryWrapper<Dataset> lqw = new LambdaQueryWrapper<>();
        lqw.eq(Dataset::getDatasetUsage, datasetUsage)
                .eq(Dataset::getDatasetStatus, DatasetStatusEnum.DONE.getStatus())
                .eq(Dataset::getDatasetType, model.getModelType());
        return datasetConvert.po2TrainingDatasetQueryVOList(datasetService.list(lqw));

    }

    @Override
    public ModelConfig getModelConfig(Integer modelId) {
        // 简单理解：根据模型id查询模型config
        // 模型必须存在
        Model model = BriupAssert.requireNotNull(
                modelService,
                Model::getId,
                modelId,
                ResultCodeEnum.DATA_NOT_EXIST
        );

        // 模型的训练状态必须是已完成
        BriupAssert.requireEqual(
                model.getTrainingStatus(),
                TrainingStatusEnum.DONE.getStatus(),
                ResultCodeEnum.PARAM_IS_ERROR);

        // 查询
        LambdaQueryWrapper<ModelConfig> lqw = new LambdaQueryWrapper<>();
        lqw.eq(ModelConfig::getModelId, modelId);
        // 一个模型训练只能有一个配置
        return modelConfigService.getOne(lqw);
    }

    @Override
    @Transactional
    public void operationModel(ModelOperationDTO modelOperationDTO){
        // 获取参数中的数据
        Integer modelId = modelOperationDTO.getModelId();
        Integer operationType = modelOperationDTO.getOperationType();
        List<Integer> datasetIds = modelOperationDTO.getDatasetIds();
        // 验证模型必须存在
        Model model = BriupAssert.requireNotNull(
                modelService,
                Model::getId,
                modelId,
                ResultCodeEnum.DATA_NOT_EXIST
        );
        // 获取模型版本
        Integer modelVersion = model.getLastModelVersion();
        // 查询训练
        LambdaQueryWrapper<Training> lqw = new LambdaQueryWrapper<>();
        lqw.eq(Training::getModelId, modelId).eq(Training::getModelVersion, modelVersion);
        Training training = trainingService.getOne(lqw);
        // 分出不同的操作类型
        if(ObjectUtil.equal(operationType, ModelOperationTypeEnum.INIT.getType())){
            // 初始化训练
            // 模型必须是未训练的状态
            BriupAssert.requireEqual(
                    model.getTrainingStatus(),
                    TrainingStatusEnum.NO_TRAINING.getStatus(),
                    ResultCodeEnum.PARAM_IS_ERROR);
            // 模型的训练记录必须为空
            BriupAssert.requireNull(training, ResultCodeEnum.DATA_ALREADY_EXIST);
            // 保存配置信息到数据库中
            modelConfigService.save(modelConfigConvert.modelOperationDTO2Po(modelOperationDTO));
            // 保存训练的相关记录到数据库中
            // 获取代理类对象
            Integer trainingId = ((ModelOperationServiceImpl) AopContext.currentProxy())
                    .saveTrainingRecord(modelId, modelVersion + 1, datasetIds);
            // 合并初始化的数据集
            Map<String, List<String>> trainingData = mergeDatasets(datasetIds);
            // 向mq中发送初始化训练的消息ModelInitTrainMessage
            ModelInitTrainMessage modelInitTrainMessage = modelConvert.modelOperationDTO2InitMessage(modelOperationDTO);
            modelInitTrainMessage.setData(trainingData);
            modelInitTrainMessage.setModelVersion(modelVersion+1);
            modelInitTrainMessage.setTrainingId(trainingId);
            rabbitTemplate.convertAndSend(
                    initExchange,
                    initRoutingKey ,
                    JSON.toJSONString(modelInitTrainMessage)
                    );
            // 修改模型状态为训练中
            model.setTrainingStatus(TrainingStatusEnum.TRAINING.getStatus());
            modelService.updateById(model);

        }else if(ObjectUtil.equal(operationType, ModelOperationTypeEnum.OPTIMIZE.getType())){
            // 优化训练
            // 模型必须是训练完成的状态
            BriupAssert.requireEqual(
                    model.getTrainingStatus(),
                    TrainingStatusEnum.DONE.getStatus(),
                    ResultCodeEnum.PARAM_IS_ERROR);
            // 保存训练揭露
            Integer trainingId = ((ModelOperationServiceImpl) AopContext.currentProxy())
                    .saveTrainingRecord(modelId, modelVersion + 1, datasetIds);
            // 合并初始化的数据集
            Map<String, List<String>> trainingData = mergeDatasets(datasetIds);
            // 向mq中发送优化训练的消息ModelOptimizeTrainMessage
            ModelOptimizeTrainMessage modelOptimizeTrainMessage = modelConvert.modelOperationDTO2OptimizeMessage(modelOperationDTO);
            modelOptimizeTrainMessage.setData(trainingData);
            modelOptimizeTrainMessage.setTrainingId(trainingId);
            modelOptimizeTrainMessage.setModelVersion(modelVersion+1);
            modelOptimizeTrainMessage.setOldModelPath(training.getModelFileAddr());
            rabbitTemplate.convertAndSend(
                    optimizeExchange,
                    optimizeRoutingKey ,
                    JSON.toJSONString(modelOptimizeTrainMessage));
            // 修改模型的状态为优化中
            model.setTrainingStatus(TrainingStatusEnum.OPTIMIZING.getStatus());
            modelService.updateById(model);

        }else if(ObjectUtil.equal(operationType, ModelOperationTypeEnum.EVALUATE.getType())){
            // 模型评估
            // 模型必须是训练完成
            BriupAssert.requireEqual(
                    model.getTrainingStatus(),
                    TrainingStatusEnum.DONE.getStatus(),
                    ResultCodeEnum.PARAM_IS_ERROR);
            // 同一个模型版本下，不能使用相同的数据集
            LambdaQueryWrapper<Evaluate> evaluateLambdaQueryWrapper = new LambdaQueryWrapper<>();
            evaluateLambdaQueryWrapper.eq(Evaluate::getTrainingId, training.getId());
            List<Integer> dataSetList = evaluateService.list(evaluateLambdaQueryWrapper).stream().map(Evaluate::getDatasetId).toList();
            BriupAssert.requireNotIn(
                    datasetIds.get(0),
                    dataSetList,
                    ResultCodeEnum.MODEL_CAN_NOT_EVALUATE);
            // 保存评估-数据集的记录
            Evaluate evaluate = new Evaluate();
            evaluate.setTrainingId(training.getId());
            evaluate.setDatasetId(datasetIds.get(0));
            evaluateService.save(evaluate);
            // 获取评估的数据集，需要合并
            Map<String, List<String>> evaluateData = mergeDatasets(datasetIds);
            // 向mq中发送评估训练的消息ModelEvaluateMessage
            ModelEvaluateMessage modelEvaluateMessage = new ModelEvaluateMessage();
            modelEvaluateMessage.setModelId(modelId);
            modelEvaluateMessage.setEvaluateId(evaluate.getId());
            modelEvaluateMessage.setData(evaluateData);
            modelEvaluateMessage.setModelFileAddr(training.getModelFileAddr());
            rabbitTemplate.convertAndSend(
                    evaluateExchange,
                    evaluateRoutingKey ,
                    JSON.toJSONString(modelEvaluateMessage)
            );
            // 更新模型状态
            model.setTrainingStatus(TrainingStatusEnum.EVALUATING.getStatus());
            modelService.updateById(model);

        }else{
            BriupAssert.throwException(ResultCodeEnum.PARAM_IS_ERROR);
        }
    }

    @Override
    public void releaseModelOrNot(Integer modelId, Integer modelStatus) {
        // 参数校验
        // 模型必须存在
        Model model = BriupAssert.requireNotNull(
                modelService,
                Model::getId,
                modelId,
                ResultCodeEnum.DATA_NOT_EXIST);
        // model的状态必须一致
        BriupAssert.requireEqual(
                model.getModelStatus(),
                modelStatus,
                ResultCodeEnum.PARAM_IS_ERROR);
        // 模型的准确率不能低于阈值
        BriupAssert.requireTrue(
                model.getAccuracyRate()>ModelConstant.MODEL_RELEASE_THRESHOLD,
                ResultCodeEnum.MODEL_CAN_NOT_RELEASE);

        // 发送模型发布的消息给mq
        // 构建消息对象
        ModelReleaseMessage modelReleaseMessage = new ModelReleaseMessage();
        // 根据版本和modelId查询训练
        LambdaQueryWrapper<Training> lqw =  new LambdaQueryWrapper<>();
        lqw.eq(Training::getModelId, modelId).eq(Training::getModelVersion, model.getLastModelVersion());
        Training training = trainingService.getOne(lqw);
        modelReleaseMessage.setModelStatus(modelStatus);
        modelReleaseMessage.setModelId(training.getModelId());
        modelReleaseMessage.setModelFileAddr(training.getModelFileAddr());
        modelReleaseMessage.setUserId(SecurityUtil.getUserId());
        rabbitTemplate.convertAndSend(
                releaseExchange,
                releaseRoutingKey,
                JSON.toJSONString(modelReleaseMessage)
        );
    }

    @Override
    public List<ReleaseModelVO> getReleaseModel() {
        // 查询所有的已发布的模型id
        return releaseService.list().stream()
                .map(s-> modelConvert.po2ReleaseModelVO(modelService.getById(s.getModelId()))
        ).toList();
    }

    @SneakyThrows // 隐藏异常抛出
    @Override
    public String identify(Integer modelId, MultipartFile file) {
        // 识别的意思就是发送一个文件给python端，之前的通信都是mq。这次使用的是http
        // 意思就是将当前的程序作为客户端，发送带文件的请求到python端
        // 要在发布表中根据模型id查询发布的url
        LambdaQueryWrapper<Release> lqw = new LambdaQueryWrapper<>();
        lqw.eq(Release::getModelId, modelId);
        Release release = releaseService.getOne(lqw);
        String url = release.getModelUrl();
        // 生成发送文件的http请求，设置请求头和请求参数
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.MULTIPART_FORM_DATA);
        MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
        body.add("file", new ByteArrayResource(file.getBytes()){
            @Override
            public String getFilename() {
                return file.getOriginalFilename();
            }
        });
        // 创建请求实体，使用restTemplate发送请求
        HttpEntity<MultiValueMap<String, Object>> requestEntity = new HttpEntity<>(body, headers);
        ResponseEntity<String> response = restTemplate.exchange(
                url,
                HttpMethod.POST,
                requestEntity,
                String.class
        );

        return response.getBody();
    }


    // ============= 封装方法 ==================

    /**
     * 保存训练相关的数据到数据库中（训练记录、绑定训练数据集）
     * @param modelId
     * @param modelVersion
     * @param datasetIds
     * @return
     */
    @Transactional
    public Integer saveTrainingRecord(Integer modelId, Integer modelVersion, List<Integer> datasetIds){
        // 新建训练对象进行保存
        Training training = new Training();
        training.setModelId(modelVersion);
        training.setModelVersion(modelVersion);
        training.setCreateTime(new Date());
        training.setCreateBy(SecurityUtil.getUserId());
        trainingService.save(training);

        // 新建训练-数据集对象进行保存
        Integer  trainingId= training.getId();
        // 输入数据集id，输出训练数据集对象
        List<TrainingDataset> list = datasetIds.stream().map(datasetId -> {
            TrainingDataset trainingDataset = new TrainingDataset();
            trainingDataset.setTrainingId(trainingId);
            trainingDataset.setDatasetId(datasetId);
            return trainingDataset;
        }).toList();
        trainingDatasetService.saveBatch(list);
        return trainingId;
    }
    /**
     * 合并数据集
     * @param datasetIds 数据集id集合
     * @return 根据分类的到合并后的数据集（key是分类名称，value是该分类下所有图片的url集合）
     */
    private Map<String, List<String>> mergeDatasets(List<Integer> datasetIds) {
        Map<String, List<String>> map = new HashMap<>();
        // 查询该数据集下的所有分类
        datasetIds.forEach(datasetId-> {
                    LambdaQueryWrapper<Classify> classifyLambdaQueryWrapper = new LambdaQueryWrapper<>();
                    classifyLambdaQueryWrapper.eq(Classify::getDatasetId, datasetId);
                    List<Classify> classifies = classifyService.list(classifyLambdaQueryWrapper);
                    // 继续查询每个分类下的图片
                    classifies.forEach(classify -> {
                        String classifyName = classify.getClassifyName();
                        Integer classifyId = classify.getId();
                        LambdaQueryWrapper<Entity> entityLambdaQueryWrapper = new LambdaQueryWrapper<>();
                        entityLambdaQueryWrapper.eq(Entity::getClassifyId, classify.getId());
                        List<Entity> entities = entityService.list(entityLambdaQueryWrapper);
                        List<String> entityUrls = new ArrayList<>(entities.stream().map(entity ->
                            CommonConstant.createEntityPath(
                                    this.nginxFilePath,
                                    datasetId,
                                    classifyName,
                                    entity.getEntityUrl())
                        ).toList());
                        // 判断map中分类名称是否已经存在
                        if(map.containsKey(classifyName)){
                            // 存在就直接拼接
                            map.get(classifyName).addAll(entityUrls);
                        }else {
                            // 不存在就存个新的
                            map.put(classifyName, entityUrls);
                        }
                    });
                }
        );

        return map;
    }
}